#!/usr/bin/python

import matplotlib.pyplot as plt
from prep_terrain_data import makeTerrainData
from class_vis import prettyPicture

from sklearn.neighbors import KNeighborsClassifier

features_train, labels_train, features_test, labels_test = makeTerrainData()


### the training data (features_train, labels_train) have both "fast" and "slow"
### points mixed together--separate them so we can give them different colors
### in the scatterplot and identify them visually
grade_fast = [features_train[ii][0] for ii in range(0, len(features_train)) if labels_train[ii]==0]
bumpy_fast = [features_train[ii][1] for ii in range(0, len(features_train)) if labels_train[ii]==0]
grade_slow = [features_train[ii][0] for ii in range(0, len(features_train)) if labels_train[ii]==1]
bumpy_slow = [features_train[ii][1] for ii in range(0, len(features_train)) if labels_train[ii]==1]


#### initial visualization
# plt.xlim(0.0, 1.0)
# plt.ylim(0.0, 1.0)
# plt.scatter(bumpy_fast, grade_fast, color = "b", label="fast")
# plt.scatter(grade_slow, bumpy_slow, color = "r", label="slow")
# plt.legend()
# plt.xlabel("bumpiness")
# plt.ylabel("grade")
# plt.show()
################################################################################


### your code here!  name your classifier object clf if you want the 
### visualization code (prettyPicture) to show you the decision boundary

def try_k_nearest():
    test_k = [5, 20, 50]
    for k in test_k:
        clf = KNeighborsClassifier(n_neighbors=k)
        clf.fit(features_train, labels_train)
        acc = clf.score(features_test, labels_test)
        try:
            title = "K-nearest neighbors when k=%i (accuracy=%0.2f%%)" % (k, acc * 100)
            name = "knearest%i.png" % k
            prettyPicture(clf, features_test, labels_test, title=title, saveto=name)
        except Exception as e:
            print(e)

def find_sweet_spot(k, v):
    for i in range(k-v, k+v):
        clf = KNeighborsClassifier(n_neighbors=i)
        clf.fit(features_train, labels_train)
        accuracy = clf.score(features_test, labels_test)
        print("k=%i, accuracy=%0.4f" % (i, accuracy))

if __name__ == "__main__":
    # try_k_nearest()
    find_sweet_spot(20,3)

